#!/usr/bin/env python3
"""
ATZLinux Store Daemon - Root权限的apt执行服务
负责执行需要root权限的apt相关命令
"""
import os
import sys
import json
import logging
import socket
import threading
import time
from subprocess import Popen, PIPE, STDOUT
from logging.handlers import TimedRotatingFileHandler

# 配置日志
atzdir = '/var/lib/atzlinux-store'
if not os.path.exists(atzdir):
    os.makedirs(atzdir, mode=0o755)

fmt = logging.Formatter('%(asctime)s - %(levelname)s %(message)s')
console = logging.StreamHandler()
console.setFormatter(fmt)

logfile = f'{atzdir}/daemon.log'
filelog = TimedRotatingFileHandler(logfile,
    when='D', backupCount=10, encoding='utf-8')
filelog.setFormatter(fmt)

logging.getLogger('').addHandler(console)
logging.getLogger('').addHandler(filelog)
logging.getLogger('').setLevel(logging.INFO)

# Unix Domain Socket 路径
SOCKET_PATH = '/var/run/atzstore-daemon.sock'

class AptDaemon:
    """APT操作守护进程"""
    
    def __init__(self):
        self.socket = None
        self.running = False
        self.clients = []
        
    def start(self):
        """启动守护进程"""
        logging.info("ATZStore Daemon starting...")
        
        # 检查是否以root身份运行
        if os.geteuid() != 0:
            logging.error("Daemon must be run as root")
            sys.exit(1)
            
        # 清理旧的socket文件
        if os.path.exists(SOCKET_PATH):
            os.unlink(SOCKET_PATH)
            
        # 创建Unix Domain Socket
        self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.socket.bind(SOCKET_PATH)
        os.chmod(SOCKET_PATH, 0o666)  # 允许普通用户连接
        self.socket.listen(5)
        
        logging.info(f"Daemon listening on {SOCKET_PATH}")
        self.running = True
        
        try:
            while self.running:
                client_socket, _ = self.socket.accept()
                logging.info("New client connected")
                
                # 为每个客户端创建处理线程
                client_thread = threading.Thread(
                    target=self.handle_client,
                    args=(client_socket,)
                )
                client_thread.daemon = True
                client_thread.start()
                
        except KeyboardInterrupt:
            logging.info("Received interrupt signal")
        finally:
            self.stop()
            
    def stop(self):
        """停止守护进程"""
        logging.info("Stopping daemon...")
        self.running = False
        
        if self.socket:
            self.socket.close()
            
        if os.path.exists(SOCKET_PATH):
            os.unlink(SOCKET_PATH)
            
    def handle_client(self, client_socket):
        """处理客户端连接"""
        try:
            while self.running:
                # 接收消息
                data = client_socket.recv(4096)
                if not data:
                    break
                    
                try:
                    message = json.loads(data.decode('utf-8'))
                    logging.info(f"Received message: {message}")
                    
                    # 处理消息
                    response = self.process_message(message, client_socket)
                    
                    # 发送响应
                    if response:
                        client_socket.send(json.dumps(response).encode('utf-8'))
                        
                except json.JSONDecodeError as e:
                    logging.error(f"Invalid JSON message: {e}")
                    error_response = {
                        "type": "error",
                        "message": "Invalid JSON format"
                    }
                    client_socket.send(json.dumps(error_response).encode('utf-8'))
                    
        except Exception as e:
            logging.error(f"Error handling client: {e}")
        finally:
            client_socket.close()
            logging.info("Client disconnected")
            
    def process_message(self, message, client_socket):
        """处理客户端消息"""
        msg_type = message.get('type')
        
        if msg_type == 'ping':
            return {"type": "pong"}
            
        elif msg_type == 'apt_command':
            command = message.get('command', '')
            if self.is_safe_command(command):
                # 在后台线程中执行命令
                cmd_thread = threading.Thread(
                    target=self.execute_apt_command,
                    args=(command, client_socket)
                )
                cmd_thread.daemon = True
                cmd_thread.start()
                
                return {"type": "command_started", "command": command}
            else:
                return {"type": "error", "message": f"Unsafe command: {command}"}
                
        elif msg_type == 'get_installed_packages':
            # 获取已安装软件包列表
            return self.get_installed_packages()
            
        else:
            return {"type": "error", "message": f"Unknown message type: {msg_type}"}
            
    def is_safe_command(self, command):
        """检查命令是否安全"""
        # 只允许特定的apt命令
        safe_commands = [
            'apt update',
            'apt install',
            'apt remove',
            'apt purge', 
            'apt upgrade',
            'apt list --installed'
        ]
        
        # 检查命令是否以安全前缀开始
        for safe_cmd in safe_commands:
            if command.strip().startswith(safe_cmd):
                return True
                
        return False
        
    def execute_apt_command(self, command, client_socket):
        """执行apt命令"""
        logging.info(f"Executing command: {command}")
        
        # 设置环境变量
        env = os.environ.copy()
        env['DEBIAN_FRONTEND'] = 'noninteractive'
        env['DEBCONF_NONINTERACTIVE_SEEN'] = 'true'
        
        try:
            # 发送开始执行消息
            start_msg = {
                "type": "command_output",
                "command": command,
                "status": "started"
            }
            client_socket.send(json.dumps(start_msg).encode('utf-8'))
            
            # 执行命令
            process = Popen(
                command,
                stdout=PIPE,
                stderr=STDOUT,
                shell=True,
                env=env,
                universal_newlines=True
            )
            
            # 实时发送输出
            for line in iter(process.stdout.readline, ''):
                if line.strip() and 'stable CLI interface' not in line:
                    output_msg = {
                        "type": "command_output",
                        "command": command,
                        "output": line.strip(),
                        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
                    }
                    try:
                        client_socket.send(json.dumps(output_msg).encode('utf-8'))
                    except:
                        # 客户端可能已断开连接
                        break
                        
            process.stdout.close()
            return_code = process.wait()
            
            # 发送完成消息
            complete_msg = {
                "type": "command_complete",
                "command": command,
                "return_code": return_code,
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
            }
            client_socket.send(json.dumps(complete_msg).encode('utf-8'))
            
            logging.info(f"Command completed with return code: {return_code}")
            
        except Exception as e:
            logging.error(f"Error executing command: {e}")
            error_msg = {
                "type": "command_error",
                "command": command,
                "error": str(e)
            }
            try:
                client_socket.send(json.dumps(error_msg).encode('utf-8'))
            except:
                pass
                
    def get_installed_packages(self):
        """获取已安装软件包列表"""
        try:
            process = Popen(
                "apt list --installed",
                stdout=PIPE,
                stderr=STDOUT,
                shell=True,
                universal_newlines=True
            )
            
            packages = {}
            for line in process.stdout:
                parts = line.strip().split('/')
                if len(parts) >= 2:
                    pkg_name = parts[0]
                    # 解析版本和状态信息
                    info_parts = line.split()
                    if len(info_parts) >= 3:
                        version = info_parts[1]
                        status = info_parts[2] if len(info_parts) > 2 else ""
                        packages[pkg_name] = f"{version}{status}"
                        
            process.wait()
            
            return {
                "type": "installed_packages",
                "packages": packages
            }
            
        except Exception as e:
            logging.error(f"Error getting installed packages: {e}")
            return {
                "type": "error", 
                "message": f"Failed to get installed packages: {e}"
            }

def main():
    """主函数"""
    daemon = AptDaemon()
    
    try:
        daemon.start()
    except Exception as e:
        logging.error(f"Daemon error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()